{- Axel Lyckberg, lyckberg@student.chalmers.se
 - Alexander Spalding, spalding.a@gmail.com
 - Chalmers Univeristy of Technology 2009
 - http://code.google.com/p/fopl09deforesting/
-}

module Deforestation where

import Control.Monad.State
import Environment
import Absfgrammar
import ErrM
import System ( exitFailure )

import Printfgrammar (printTree)

-- Constant True if debug is to be printed, False ow.
-- TODO: Give this as a flag to Main.hs instead
debug :: Bool
debug = True

-- Debug printing function.
printDebug :: Int -> Exp -> Exp -> TState ()
printDebug rnum exp1 exp2 = do
    case debug of
        True -> do
            liftIO $ putStrLn $ 
                "\n=========== RULE " ++ (show rnum) ++ " ============="
            liftIO $ putStrLn $ show exp1
            liftIO $ putStrLn $ "========= REPLACED WITH ========="
            liftIO $ putStrLn $ show exp2
            liftIO $ putStrLn $ "========= Prettyprint =========="
            liftIO $ putStrLn $ printTree exp2
            liftIO $ getChar
            return ()
        _ -> return ()


-- Entrypoint, construct State-monad from Environment.hs.
-- Read signature and transform program to treeless form.
deforest :: Program -> IO Program
deforest prog = do
    (tp, _) <- runStateT (deforest' prog) s0
    return tp
    where
    deforest' :: Program -> TState Program
    deforest' p = do
        readDefs p
        sig <- getSig
        transProgram p
        tp <- getProgram
        return tp

readDefs :: Program -> TState ()
readDefs (Prog defs) = mapM_ addFun defs

----------- IMPLEMENT TODO TO STATEMONAD ------------------
transProgram :: Program -> TState ()
transProgram (Prog defs) = mapM_ transDef defs

transDef :: Def -> TState ()
transDef d = case d of
    Func ids exp  -> do
        exp' <- transExp exp
        addDef (Funr ids exp')
    _  -> addDef d

transExp :: Exp -> TState Exp
transExp exp = case exp of
    -------- Rule 1 ----------
    EId id -> return exp
    EInt lit -> return exp
    -------- Rule 2 ----------
    EDiv e1 e2     -> do 
        (e1', e2') <- transBinOp e1 e2
        return (EDiv e1' e2')
    EMul e1 e2     -> do 
        (e1', e2') <- transBinOp e1 e2
        return (EMul e1' e2')
    EAdd e1 e2     -> do 
        (e1', e2') <- transBinOp e1 e2
        return (EAdd e1' e2')
    ESub e1 e2     -> do 
        (e1', e2') <- transBinOp e1 e2
        return (ESub e1' e2')
    EGre e1 e2     -> do 
        (e1', e2') <- transBinOp e1 e2
        return (EGre e1' e2')
    ELes e1 e2     -> do 
        (e1', e2') <- transBinOp e1 e2
        return (ELes e1' e2')
    EGte e1 e2     -> do 
        (e1', e2') <- transBinOp e1 e2
        return (EGte e1' e2')
    ELte e1 e2     -> do 
        (e1', e2') <- transBinOp e1 e2
        return (ELte e1' e2')
    EEqu e1 e2     -> do 
        (e1', e2') <- transBinOp e1 e2
        return (EEqu e1' e2')
    ECon id exps   -> do
        exps' <- mapM transExp exps
        return (ECon id exps')
    -------- Rule 3 ----------
    EApp (EId id) exp2 -> do -- Rule 3
        -- Replace all occurences of "the single argument?" in fun with exp2
        fun <- getFun id
        putReplacementExp exp2
        exp' <- replaceArg fun
        if exp' == (EId (Ident ""))
            then return exp
            else do
                printDebug 3 exp exp'
                -- Get new name for function
                exp'' <- transExp exp'
                return exp''
    -------- Rule 4 ----------
    ECas (EId id) mtchs -> do
        mtchs' <- mapM transMatch mtchs
        return (ECas (EId id) mtchs')
    ECas (EInt lit) mtchs -> do 
        mtchs' <- mapM transMatch mtchs
        return (ECas (EInt lit) mtchs')
    -------- Rule 5 ----------
    ECas (ECon id terms) matchs -> do
        (EMat e1 e2) <- findMatch id matchs
        -- Replace occurences of exps inside e2 with e1 
        -- let fun = (Func [(Ident "")] e2)
        putReplaceTerms terms -- The exps to find inside e2
        putReplacementExp e1 -- The 
        exp' <- replaceExps e2
        printDebug 5 exp exp' 
        return e2
    --    liftIO $ putStrLn $ show exp
    --    fail $ "RULE 5 NOT IMPLEMENTED"
    -------- Rule 6 ----------
    ECas (EApp (EId id) exp2) matchs -> do
        fun <- getFun id
        putReplacementExp exp2
        exp' <- replaceArg fun
        if exp' == (EId (Ident ""))
            then return exp
            else do
                printDebug 6 exp (ECas exp' matchs)
                exp'' <- transExp (ECas exp' matchs)
                return exp''
    -------- Rule 7 ----------
    ECas (ECas exp1 matchs1) matchs2 -> do
        newCases <- buildNewCases exp1 matchs1 matchs2
        printDebug 7 exp newCases
        exp' <- transExp newCases
        return exp'
    EApp exp1 exp2 -> 
        fail $ "\nEXP 1\n" ++ (show exp1) ++ "\nEXP 2\n" ++ (show exp2)
    _ -> fail $ "inte okej\n\n\n" ++ (show exp)

replaceExps :: Exp -> TState Exp
replaceExps exp = case exp of
    EId id   -> do
        terms <- getReplaceTerms
        case (exp `elem` terms) of
            True -> do
                exp' <- getReplacementExp
                return exp'
            _    -> return exp
    EInt lit -> do
        terms <- getReplaceTerms
        case (exp `elem` terms) of
            True -> do
                exp' <- getReplacementExp
                return exp'
            _    -> return exp
    ECon id exps  -> do
        exps' <- mapM replaceExps exps
        return (ECon id exps')
    EApp exp0 exp1  -> do
        exp0' <- replaceExps exp0
        exp1' <- replaceExps exp1
        return (EApp exp0' exp1')
    EDiv exp0 exp1  -> do
        exp0' <- replaceExps exp0
        exp1' <- replaceExps exp1
        return (EDiv exp0' exp1')
    EMul exp0 exp1  -> do
        exp0' <- replaceExps exp0
        exp1' <- replaceExps exp1
        return (EMul exp0' exp1')
    EAdd exp0 exp1  -> do
        exp0' <- replaceExps exp0
        exp1' <- replaceExps exp1
        return (EAdd exp0' exp1')
    ESub exp0 exp1  -> do
        exp0' <- replaceExps exp0
        exp1' <- replaceExps exp1
        return (ESub exp0' exp1')
    EGre exp0 exp1  -> do
        exp0' <- replaceExps exp0
        exp1' <- replaceExps exp1
        return (EGre exp0' exp1')
    ELes exp0 exp1  -> do
        exp0' <- replaceExps exp0
        exp1' <- replaceExps exp1
        return (ELes exp0' exp1')
    EGte exp0 exp1  -> do
        exp0' <- replaceExps exp0
        exp1' <- replaceExps exp1
        return (EGte exp0' exp1')
    ELte exp0 exp1  -> do
        exp0' <- replaceExps exp0
        exp1' <- replaceExps exp1
        return (ELte exp0' exp1')
    EEqu exp0 exp1  -> do
        exp0' <- replaceExps exp0
        exp1' <- replaceExps exp1
        return (EEqu exp0' exp1')
    ECas exp0 matchs -> do
        matchs' <- mapM replaceMatchs matchs
        exp0' <- replaceExps exp0
        return (ECas exp0' matchs')

replaceMatchs :: Match -> TState Match
replaceMatchs (EMat e1 e2) = do
    e1' <- replaceExps e1
    e2' <- replaceExps e2
    return (EMat e1' e2')


transBinOp :: Exp -> Exp -> TState (Exp, Exp)
transBinOp e1 e2 = do
    e1 <- transExp e1
    e2 <- transExp e2
    return (e1, e2)

transMatch :: Match -> TState Match
transMatch (EMat e1 e2) = do
    e2' <- transExp e2
    return (EMat e1 e2')

buildNewCases :: Exp -> [Match] -> [Match] -> TState Exp
buildNewCases e ms m2s = 
    return (ECas e [(EMat e1 (ECas e2 m2s)) | (EMat e1 e2) <- ms])

findMatch :: Ident -> [Match] -> TState Match
findMatch id [] = fail $ "Didn't find constructor " ++ (show id)
findMatch id (m:ms) = case m of
    EMat (ECon id' _) _ -> if (id == id')
                                    then return m
                                    else findMatch id ms
    _                 -> findMatch id ms

replaceArg :: Def -> TState Exp
replaceArg d = case d of
    (Funr (ida:arg:[]) exp) -> do
        putArgName arg
        exp' <- replaceInExp exp
        return exp'
    (Func (ida:arg:[]) exp) -> do
        putArgName arg
        exp' <- replaceInExp exp
        return exp'
    _ -> return (EId (Ident ""))

replaceInExp :: Exp -> TState Exp
replaceInExp exp = case exp of
    EId id' -> do
        id <- getArgName
        if id == id'
            then do
                exp' <- getReplacementExp
                return exp'
            else return (EId id')
    EInt literal  -> return exp
    ECon id exps  -> do
        exps' <- mapM replaceInExp exps
        return (ECon id exps')
    EApp exp0 exp1  -> do
        exp0' <- replaceInExp exp0
        exp1' <- replaceInExp exp1
        return (EApp exp0' exp1')
    EDiv exp0 exp1  -> do
        exp0' <- replaceInExp exp0
        exp1' <- replaceInExp exp1
        return (EDiv exp0' exp1')
    EMul exp0 exp1  -> do
        exp0' <- replaceInExp exp0
        exp1' <- replaceInExp exp1
        return (EMul exp0' exp1')
    EAdd exp0 exp1  -> do
        exp0' <- replaceInExp exp0
        exp1' <- replaceInExp exp1
        return (EAdd exp0' exp1')
    ESub exp0 exp1  -> do
        exp0' <- replaceInExp exp0
        exp1' <- replaceInExp exp1
        return (ESub exp0' exp1')
    EGre exp0 exp1  -> do
        exp0' <- replaceInExp exp0
        exp1' <- replaceInExp exp1
        return (EGre exp0' exp1')
    ELes exp0 exp1  -> do
        exp0' <- replaceInExp exp0
        exp1' <- replaceInExp exp1
        return (ELes exp0' exp1')
    EGte exp0 exp1  -> do
        exp0' <- replaceInExp exp0
        exp1' <- replaceInExp exp1
        return (EGte exp0' exp1')
    ELte exp0 exp1  -> do
        exp0' <- replaceInExp exp0
        exp1' <- replaceInExp exp1
        return (ELte exp0' exp1')
    EEqu exp0 exp1  -> do
        exp0' <- replaceInExp exp0
        exp1' <- replaceInExp exp1
        return (EEqu exp0' exp1')
    ECas exp0 matchs -> do
        matchs' <- mapM replaceInMatch matchs
        exp0' <- replaceInExp exp0
        return (ECas exp0' matchs')

replaceInMatch :: Match -> TState Match
replaceInMatch (EMat e1 e2) = do
    e1' <- replaceInExp e1
    e2' <- replaceInExp e2
    return (EMat e1' e2')
    

----------------------------------------------------------------

type Result = Err String

failure :: Show a => a -> Result
failure x = Bad $ "Undefined case: " ++ show x

transIdent :: Ident -> Result
transIdent x = case x of
  Ident str  -> failure x
{-
fail :: String -> TState ()
fail s = do
    liftIO $ putStrLn $ s
    liftIO $ exitFailure
    -}

{-
transProgram :: Program -> Result
transProgram x = case x of
  Prog defs  -> failure x


transDef :: Def -> Result
transDef x = case x of
  Func ids exp  -> failure x
  Funr ids exp  -> failure x
  Funp str exp  -> failure x


transExp :: Exp -> Result
transExp x = case x of
  EId id  -> failure x
  EInt literal  -> failure x
  ECon id exps  -> failure x
  EApp exp0 exp  -> failure x
  EDiv exp0 exp  -> failure x
  EMul exp0 exp  -> failure x
  EAdd exp0 exp  -> failure x
  ESub exp0 exp  -> failure x
  EGre exp0 exp  -> failure x
  ELes exp0 exp  -> failure x
  EGte exp0 exp  -> failure x
  ELte exp0 exp  -> failure x
  EEqu exp0 exp  -> failure x
  ECas exp matchs  -> failure x


transMatch :: Match -> Result
transMatch x = case x of
  EMat exp0 exp  -> failure x


transLiteral :: Literal -> Result
transLiteral x = case x of
  Lint n  -> failure x
  Ldbl d  -> failure x
  Ltru  -> failure x
  Lfal  -> failure x


-}
